{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JNSdeDvgpdng"
      },
      "source": [
        "# Notebook Setup\n",
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/PacktPublishing/Generative-AI-Integration-Patterns-1E/blob/main/Chapter08/Integration_pattern_Real_Time_retrieval_augmented_generation.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "nZzaJzUUoRbW",
        "outputId": "d3401be0-b0f3-4b14-ff74-b7510923eb56"
      },
      "outputs": [],
      "source": [
        "#Install dependencies\n",
        "\n",
        "!pip install --upgrade google-cloud-aiplatform\n",
        "!pip install --upgrade langchain_community langchain_google_vertexai langchain_chroma unstructured[pdf]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "euBE2B3DpHSQ"
      },
      "outputs": [],
      "source": [
        "#Authenticate\n",
        "from google.colab import auth as google_auth\n",
        "google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IHRIcyzsoT6J"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import json\n",
        "\n",
        "#VertexAI\n",
        "import vertexai\n",
        "from vertexai.generative_models import GenerativeModel, Part, FinishReason\n",
        "import vertexai.preview.generative_models as generative_models\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "# Langchain\n",
        "from langchain_community.document_loaders import TextLoader, UnstructuredPDFLoader\n",
        "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
        "from unstructured.cleaners.core import clean_extra_whitespace\n",
        "from langchain_google_vertexai import VertexAIEmbeddings\n",
        "\n",
        "\n",
        "#Markdown\n",
        "from IPython.display import display, Markdown, Latex\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qgNuihpJqgkJ"
      },
      "outputs": [],
      "source": [
        "PROJECT = \"testproject-410220\"#@param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"#@param {type:\"string\"}\n",
        "MODEL = \"gemini-1.5-flash-001\"#@param {type:\"string\"}\n",
        "EMBEDDINGS_MODEL = \"text-embedding-004\"#@param {type:\"string\"}\n",
        "MAX_RESULTS = 4#@param {type:\"number\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNQcTLKAC1kW"
      },
      "source": [
        "# Vector database initialization and ingestion\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XRI6Vj6yZUNh",
        "outputId": "34649eef-d43d-4a5b-b26b-fbb97923b64e"
      },
      "outputs": [],
      "source": [
        "#@title Dataset Download\n",
        "\n",
        "!wget https://d1io3yog0oux5.cloudfront.net/_c38ec26158c6d5493f3fce02d606a6a1/cocacolacompany/db/764/8109/file/CORRECTED+TRANSCRIPT_+The+Coca-Cola+Co.%28KO-US%29%2C+Q1+2024+Earnings+Call%2C+30-April-2024+8_30+AM+ET.pdf -O coca_cola_earnings_call_2023.pdf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EpduK1nxEZim",
        "outputId": "f6cbb15a-7e60-4664-8e4c-d4791e29d458"
      },
      "outputs": [],
      "source": [
        "# load the document and split it into chunks\n",
        "text_splitter = RecursiveCharacterTextSplitter(chunk_size=2048, chunk_overlap=200)\n",
        "loader = UnstructuredPDFLoader(\"coca_cola_earnings_call_2023.pdf\",post_processors=[clean_extra_whitespace])\n",
        "pages = loader.load_and_split(text_splitter=text_splitter)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zhpWQkD_a6-z",
        "outputId": "2594b4a1-5ca9-406b-862a-7d342b2a3623"
      },
      "outputs": [],
      "source": [
        "#Check what is in the chunks\n",
        "print(pages[1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4HcvcLheZ95-"
      },
      "outputs": [],
      "source": [
        "#Init VertexAI Platform\n",
        "aiplatform.init(project=PROJECT, location=LOCATION)\n",
        "embeddings_function = VertexAIEmbeddings(model=EMBEDDINGS_MODEL)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jeqDXic3iHeu"
      },
      "outputs": [],
      "source": [
        "#Create a local instance of ChromaDB\n",
        "from langchain_chroma import Chroma\n",
        "\n",
        "# Generate embeddings and load them into ChromDB\n",
        "db = Chroma.from_documents(pages, embeddings_function)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eL8x7QQrpAao",
        "outputId": "27f74fe4-ae90-43f1-f19a-6c136c3ad6aa"
      },
      "outputs": [],
      "source": [
        "#Test query\n",
        "\n",
        "# Test query\n",
        "query = \"Who is the call for?\"\n",
        "docs = db.similarity_search(query,k=MAX_RESULTS)\n",
        "\n",
        "# Print results\n",
        "print(docs[0].page_content)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fEaglMjspWeE"
      },
      "outputs": [],
      "source": [
        "#@title RAG Logic\n",
        "\n",
        "#In this section we define the prompt, as the task is to perform intent classification we will identify the intent by exposing the possible values to the LLM\n",
        "prompt_template = \"\"\"\n",
        "You are a helpful assistant for an online financial services company that allows users to check their balances, invest in certificates of deposit (CDs), and perform other financial transactions.\n",
        "\n",
        "Your task is to answer questions from your customers, in order to do so follow these rules:\n",
        "\n",
        "1. Carefully analyze the question you received.\n",
        "2. Carefully analyze the context provided.\n",
        "3. Answer the question using ONLY the information provided in the context, NEVER make up information\n",
        "4. Always think step by step.\n",
        "\n",
        "<context>\n",
        "{context}\n",
        "</context>\n",
        "User question: {query}\n",
        "Answer:\n",
        "\"\"\"\n",
        "\n",
        "generation_config = {\n",
        "    \"max_output_tokens\": 8192,\n",
        "    \"temperature\": 0,\n",
        "    \"top_p\": 0.95,\n",
        "}\n",
        "\n",
        "safety_settings = {\n",
        "    generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "}\n",
        "\n",
        "def get_context(query, db, number_of_results):\n",
        "  context_string = \"\"\n",
        "  docs = db.similarity_search(query,k=number_of_results)\n",
        "  for doc in docs:\n",
        "    new_context = f\"\"\"\\n---This information is contained in a document called {doc.metadata[\"source\"]} \\n\\n {doc.page_content}\\n\\n---\"\"\"\n",
        "    context_string = context_string+new_context\n",
        "  return(context_string)\n",
        "\n",
        "\n",
        "def generate(prompt):\n",
        "  model = GenerativeModel(MODEL)\n",
        "  responses = model.generate_content(\n",
        "      [prompt],\n",
        "      generation_config=generation_config,\n",
        "      safety_settings=safety_settings,\n",
        "      stream=False,\n",
        "  )\n",
        "  return(responses)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B7rAfCnbs22H"
      },
      "source": [
        "# Entry Point"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pg1e3tGBwgJ1"
      },
      "outputs": [],
      "source": [
        "#In this case we will simulate the input from a chat interface\n",
        "\n",
        "question = \"What is this call about?\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ihAHrCBas5eW"
      },
      "source": [
        "# Prompt Preprocessing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YkAR047_r6jL"
      },
      "outputs": [],
      "source": [
        "#In this step we will query the vector database with the question received, and then populate the promp template with both the question and the context\n",
        "context = get_context(question, db, MAX_RESULTS)\n",
        "prompt = prompt_template.format(query=question, context=context)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NdI8uxwR2ibW",
        "outputId": "f4b79953-9f57-4257-bc2e-1d7437e401dd"
      },
      "outputs": [],
      "source": [
        "print(prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LuJusiKDs7Au"
      },
      "source": [
        "# Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N_YWvX5lhHZV"
      },
      "outputs": [],
      "source": [
        "#This is the section where we submit the full prompt and context to the LLM\n",
        "result = generate(prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fnSrhoN1tACl"
      },
      "source": [
        "# Result Postprocessing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ggiDKZp3292w"
      },
      "outputs": [],
      "source": [
        "#In this section you can format the answer for example with markdown\n",
        "formatted_result = f\"###Question:\\n{question}\\n\\n###Answer:\\n{result.text}\\n\\n<details><summary>Context</summary>{context}</details>\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Vy0H0pptCCF"
      },
      "source": [
        "# Result Presentation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 166
        },
        "id": "Jv9szpkkzA23",
        "outputId": "5d901edf-29b6-4ba3-8e90-7541b0a73011"
      },
      "outputs": [],
      "source": [
        "display(Markdown(formatted_result))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h0hwuc_U2rAJ"
      },
      "source": [
        "# Demo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_vsqw1SAYRtc",
        "outputId": "6ab3892a-1618-4190-91c7-d681d16ed386"
      },
      "outputs": [],
      "source": [
        "#In this case we will use a Gradio interface to interact with the system\n",
        "\n",
        "#Install Gradio\n",
        "\n",
        "!pip install --upgrade gradio"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 684
        },
        "id": "P2JOzLS6cgeq",
        "outputId": "30669b12-1903-4f0b-b3a4-5092209d0737"
      },
      "outputs": [],
      "source": [
        "import gradio as gr\n",
        "\n",
        "def answer_question(query, db, number_of_results):\n",
        "  context = get_context(query, db, number_of_results)\n",
        "  answer = generate(prompt_template.format(query=query, context=context))\n",
        "  return(answer.text)\n",
        "\n",
        "def chat(message, history):\n",
        "    response = answer_question(message,db, MAX_RESULTS)\n",
        "    history.append((message, response))\n",
        "    return \"\", history\n",
        "\n",
        "\n",
        "with gr.Blocks() as demo:\n",
        "  gr.Markdown(\"Fintech Assistant\")\n",
        "  chatbot = gr.Chatbot(show_label=False)\n",
        "  message = gr.Textbox(placeholder=\"Enter your question\")\n",
        "  message.submit(chat, [message, chatbot],[message, chatbot]  )\n",
        "\n",
        "demo.launch(debug=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kouo1bzXC7or"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
